#!/bin/bash

set -e

# IMAGE 可通过第一个位置参数传入，示例：
#   ./optimized_test.sh randkin/llm-environment:esgzv2
# 或查看帮助：
#   ./optimized_test.sh -h
if [ "$#" -ge 1 ]; then
    case "$1" in
        -h|--help)
            echo "Usage: $0 [IMAGE]"
            echo "  IMAGE: 要测试的镜像，例如 randkin/llm-environment:esgz 或 randkin/llm-environment:esgzv2"
            exit 0
            ;;
        *)
            IMAGE="$1"
            ;;
    esac
else
    # 默认镜像（保持原脚本默认）
    IMAGE="randkin/llm-environment:esgz"
fi


MODEL_DIR="$HOME/models/TinyLlama-1.1B-Chat-v1.0-GGUF"
MODEL_FILE="tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
FULL_MODEL_PATH="$MODEL_DIR/$MODEL_FILE"

echo "========================================="
echo "  大模型推理容器冷启动测试 (nerdctl)"
echo "  镜像: $IMAGE"
echo "  模型: $(basename "$FULL_MODEL_PATH")"
echo "  大小: $(du -h "$FULL_MODEL_PATH" | cut -f1)"
echo "========================================="

# 检查模型文件
if [ ! -f "$FULL_MODEL_PATH" ]; then
    echo "错误: 模型文件不存在!"
    exit 1
fi

# 清除缓存以获得更准确的冷启动测量
echo "清除系统缓存..."
sync
echo 3 | sudo tee /proc/sys/vm/drop_caches > /dev/null

# 使用仓库中的共享测试脚本 `cold_start_test.py`
echo "Using shared cold_start_test.py from repo root"


# 先拉取镜像并计时
echo "测试镜像拉取时间..."
start_pull=$(date +%s)
sudo nerdctl --snapshotter=stargz pull $IMAGE
end_pull=$(date +%s)
PULL_TIME=$((end_pull - start_pull))
echo "镜像拉取耗时: ${PULL_TIME} 秒"
echo "========================================="

# 冷启动计时
echo "测试容器冷启动时间..."
start_time=$(date +%s%N)  # 使用纳秒精度
sudo nerdctl --snapshotter=stargz run --rm \
    -v "$MODEL_DIR":/model \
    -v $(pwd):/app \
    $IMAGE \
    python /app/cold_start_test.py
end_time=$(date +%s%N)
container_time=$(( (end_time - start_time) / 1000000 ))  # 转换为毫秒

echo "-----------------------------------------"
echo "【容器冷启动时间】: $(echo "scale=3; $container_time / 1000" | bc) 秒"
echo "========================================="

# 输出总时间（拉取+冷启动）
TOTAL_TIME_MS=$((PULL_TIME * 1000 + container_time))
echo "【镜像拉取+冷启动时间】: $(echo "scale=3; $TOTAL_TIME_MS / 1000" | bc) 秒"

# 不删除共享脚本